package com.bruce.ai.alibaba.graph.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;

import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import org.springframework.beans.factory.annotation.Qualifier;

/**
 * 人类节点的控制器类
 */
@RestController
@RequestMapping("/graph/human")
public class GraphHumanController {

    private static final Logger logger = LoggerFactory.getLogger(GraphHumanController.class);

    private final CompiledGraph compiledGraph;

    public GraphHumanController(@Qualifier("humanGraph") StateGraph stateGraph) throws GraphStateException {
        SaverConfig saverConfig = SaverConfig.builder().register(new MemorySaver())
                .build();
        this.compiledGraph = stateGraph
                .compile(CompileConfig.builder().saverConfig(saverConfig).interruptBefore("human_feedback").build());
    }

    /**
     * 展开人类节点的执行流
     * 
     * @param query         用户查询
     * @param expanderNumber 展开节点的数量
     * @param threadId      线程ID
     * @return 包含节点输出的事件流
     * @throws GraphRunnerException 如果在执行流中发生错误
     */
    @GetMapping(value = "/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(
            @RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？", required = false) String query,
            @RequestParam(value = "expander_number", defaultValue = "3", required = false) Integer expanderNumber,
            @RequestParam(value = "thread_id", defaultValue = "bruce_sg", required = false) String threadId)
            throws GraphRunnerException {
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("query", query);
        objectMap.put("expander_number", expanderNumber);

        return Mono.defer(() -> Mono.justOrEmpty(compiledGraph.invoke(objectMap, runnableConfig))
                .map(OverAllState::data)
                .map(data -> ServerSentEvent.builder(data.toString()).build()))
            .flux()
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

    /**
     * 恢复人类节点的执行流 需要先执行expand方法，构建了执行流，才能恢复执行流
     * 
     * @param threadId 线程ID  要和expand中的线程ID一致，否则会报错
     * @param feedBack 是否需要人类反馈
     * @return 包含节点输出的事件流
     * @throws GraphRunnerException 如果在执行流中发生错误
     */
    @GetMapping(value = "/resume", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> resume(
            @RequestParam(value = "thread_id", defaultValue = "bruce_sg", required = false) String threadId,
            @RequestParam(value = "feed_back", defaultValue = "true", required = false) boolean feedBack)
            throws GraphRunnerException {
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("feed_back", feedBack);

        return Mono.defer(() -> Mono.justOrEmpty(compiledGraph.invoke(objectMap, runnableConfig))
                .map(OverAllState::data)
                .map(data -> ServerSentEvent.builder(data.toString()).build()))
            .flux()
            .doOnCancel(() -> logger.info("Client disconnected from stream"))
            .doOnError(e -> logger.error("Error occurred during streaming", e));
    }
}
